package com.dap.dao.ha;

import java.sql.SQLException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicInteger;

import javax.sql.DataSource;


public class BackendPool {

    // 随机
    private final ThreadLocalRandom random;

    // 轮询
    private final ThreadLocal<AtomicInteger> atomicPositiveInteger = new ThreadLocal<AtomicInteger>() {
        @Override
        protected AtomicInteger initialValue() {
            return new AtomicInteger(0);
        };
    };

    // 可用的数据源集合
    private final ConcurrentHashMap<String, DataSource> aliveDataSourceMap;

    private final static BackendPool instance = new BackendPool();

    // 总的数据源数量,包括可用的和不可用的
    private int totalDatasourceCount = -1;
    // 是否使用轮询
    private ThreadLocal<Boolean> isRoundRobin = new ThreadLocal<Boolean>() {
        @Override
        protected Boolean initialValue() {
            return Boolean.FALSE;
        }
    };


    private BackendPool() {
        super();
        aliveDataSourceMap = new ConcurrentHashMap<String, DataSource>();
        this.random = ThreadLocalRandom.current();
    }


    public static BackendPool getInstance() {
        return instance;
    }


    public void putDataSouce(String dsName, DataSource ds) {
        aliveDataSourceMap.putIfAbsent(dsName, ds);
    }


    public void removeDataSouce(String dsName) {
        aliveDataSourceMap.remove(dsName);
    }


    /** 获取可以用的数据源集合 */
    ConcurrentHashMap<String, DataSource> getAliveDatasourceMap() {
        return aliveDataSourceMap;
    }


    /** 设置数据源总数 */
    void setTotalDatasourceCount(int totalDatasourceCount) {
        this.totalDatasourceCount = totalDatasourceCount;
    }


    /** 数据源总数，包括不可用的 */
    public int getTotalDatasourceCount() {
        return totalDatasourceCount;
    }


    /** 可用的数据源数量 */
    public int getAliveDatasourceCount() {
        return aliveDataSourceMap.size();
    }


    /** 是否多个数据源 */
    public boolean isMutilDatasource() {
        return totalDatasourceCount != -1;
    }


    /** 设置使用轮询策略 */
    public void setRoundRobin(boolean isRoundRobin) {
        this.isRoundRobin.set(isRoundRobin);
    }


    public DataSource getAliveDatasource() throws SQLException {
        int size = aliveDataSourceMap.size();
        if (size == 0) {
            throw new SQLException("All dataSource had down ! Please try again later.");
        }

        if (isRoundRobin.get()) {
            int cur = atomicPositiveInteger.get().getAndIncrement();
            // 达到最大的数据源数量，从头开始
            if (cur >= totalDatasourceCount) {
                atomicPositiveInteger.get().set(0);
                cur = atomicPositiveInteger.get().getAndIncrement();
            }
            return (DataSource) aliveDataSourceMap.values().toArray()[cur % size];
        } else
            return (DataSource) aliveDataSourceMap.values().toArray()[random.nextInt(0, size)];
    }

}
